package com.lizemin.springaspect.interceptor;

import javax.servlet.ServletOutputStream;
import javax.servlet.WriteListener;
import javax.servlet.http.HttpServletResponse;
import javax.servlet.http.HttpServletResponseWrapper;
import java.io.*;

/**
 * @author lzm
 * @date 2025/3/31 21:37
 * @description
 */
public class RepeatableResponseWrapper extends HttpServletResponseWrapper {

    /**
     * 用户缓存响应体
     */
    private ByteArrayOutputStream content = new ByteArrayOutputStream();

    /**
     * 框架会调用getOutputStream方法拿到这个流，然后调用这个流的write方法将数据写到输出流中
     */
    private ServletOutputStream outputStream;

    /**
     * 缓存PrintWriter
     */
    private PrintWriter writer;

    public RepeatableResponseWrapper(HttpServletResponse response) {
        super(response);
    }

    @Override
    public ServletOutputStream getOutputStream() {
        if (this.outputStream == null) {
            getServletOutputStream();
        }
        return this.outputStream;
    }

    // 缓存响应体
    private void getServletOutputStream() {
        this.outputStream = new ServletOutputStream() {
            @Override
            public boolean isReady() {
                return false;
            }

            @Override
            public void setWriteListener(WriteListener writeListener) {

            }

            @Override
            public void write(int b) throws IOException {
                // 往输出流中写入响应体
                RepeatableResponseWrapper.super.getOutputStream().write(b);
                // 缓存响应体
                content.write(b);
            }
        };
    }

    @Override
    public PrintWriter getWriter() throws IOException {
        if (writer == null) {
            writer = new PrintWriter(new OutputStreamWriter(this.outputStream, getCharacterEncoding()), true);
        }
        return writer;
    }

    public String getResponseBody() throws UnsupportedEncodingException {
        return this.content.toString(getCharacterEncoding());
    }

}
